#if !defined(UNICODE) || !defined(_UNICODE)
#define UNICODE
#define _UNICODE
#endif

#ifndef _WIN32_WINNT
#define _WIN32_WINNT 0x0501 // XP is the minimum
#elif (_WIN32_WINNT < 0x0501)
#error "Multi console requires XP or above"
#endif

#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif

#include <windows.h>
#include <iostream>
#include <algorithm>
#include <shlwapi.h>
#include <sstream>
#include <boost/shared_ptr.hpp>
#include "Console.h"
#include "IOStream.h"
#include "ConData.h"
#include "ChildProcessRecorder.h"

#pragma comment(lib, "rpcrt4.lib")
#pragma comment(lib, "shlwapi.lib")
#pragma comment(lib, "ole32.lib")
#pragma comment(lib, "kernel32.lib")

namespace 
{
	TCHAR consoleChildPath[MAX_PATH];
	const TCHAR* consoleChildName = TEXT("ConsoleChild.exe");
}

namespace Con
{
	namespace detail
	{
		// Keeps track of the alive children so we can bail out gracefully
		// if this process crashes
		ChildProcessRecorder recorder;

		// creates a unique pipe name by tacking a guid onto
		// the \\.\pipe prefix
		std::wstring CreateUniquePipeName()
		{
			std::wostringstream nameConstructor;
			HANDLE hCheck = INVALID_HANDLE_VALUE;
			do
			{
				if(hCheck != INVALID_HANDLE_VALUE)
				{
					CloseHandle(hCheck);
				}
				nameConstructor.str(L"");
				nameConstructor << L"\\\\.\\pipe\\";
				GUID guid;
				UuidCreateSequential(&guid);
				OLECHAR guidString[50] = {0};
				StringFromGUID2(guid, guidString, 50);
				nameConstructor << guidString;
			}
			// see if the generated name resolves to a currently available pipe
			// if it does generate a new one
			while( 
				(hCheck = CreateFileW(
					nameConstructor.str().c_str(),
					READ_CONTROL,
					0, 
					NULL, 
					OPEN_EXISTING, 
					FILE_ATTRIBUTE_NORMAL, 
					NULL))
				!= INVALID_HANDLE_VALUE);
			return nameConstructor.str();
		}

		// ensure only handle arrays of size 2 are passed in
		// by passing array references in
		bool CreatePipeHandles(HANDLE (&stdinHandles)[NUM_HANDLES], HANDLE (&stdoutHandles)[NUM_HANDLES])
		{
			bool bRet = FALSE;
			std::wstring pipeName = CreateUniquePipeName();

			SECURITY_ATTRIBUTES sa = {sizeof(sa), 0, TRUE};

			// create a named pipe so we can specify overlapped IO which is used in the child
			// processes to wait on the pipe and the exit event at the same time
			stdoutHandles[PARENT] = CreateNamedPipeW(pipeName.c_str(),
				PIPE_ACCESS_OUTBOUND | FILE_FLAG_OVERLAPPED,
				PIPE_TYPE_BYTE | PIPE_READMODE_BYTE | PIPE_WAIT,
				1,
				512,
				512,
				NMPWAIT_USE_DEFAULT_WAIT,
				NULL);

			// this might look counter-intuitive that stdout is opened for read access
			// but it is correct as the text to be written to the child's stdout 
			// will be written by the parent from the handle above.
			// The child process will read that text from this handle
			// before writing to its' screen buffer
			stdoutHandles[CHILD] = CreateFileW(
				pipeName.c_str(),
				GENERIC_READ | FILE_WRITE_ATTRIBUTES,
				0,
				&sa,
				OPEN_EXISTING,
				FILE_FLAG_OVERLAPPED | FILE_ATTRIBUTE_NORMAL,
				NULL);

			if((stdoutHandles[PARENT] != INVALID_HANDLE_VALUE) && (stdoutHandles[CHILD] != INVALID_HANDLE_VALUE))
			{
				// now the stdin handles
				// we have to use CreateNamedPipe here to create a message mode pipe
				// This is so we can detect EOF (CTRL+Z) from the child's console input buffer
				pipeName = CreateUniquePipeName();
				stdinHandles[PARENT] = CreateNamedPipeW(pipeName.c_str(),
					PIPE_ACCESS_INBOUND,
					PIPE_TYPE_MESSAGE | PIPE_READMODE_BYTE | PIPE_WAIT,
					1,
					512,
					512,
					NMPWAIT_USE_DEFAULT_WAIT,
					NULL);

				// and the other end for the child
				stdinHandles[CHILD] = CreateFileW(
					pipeName.c_str(),
					GENERIC_WRITE | FILE_READ_ATTRIBUTES,
					0,
					&sa,
					OPEN_EXISTING,
					FILE_ATTRIBUTE_NORMAL,
					NULL);

				bRet = ((stdinHandles[PARENT] != INVALID_HANDLE_VALUE) && (stdinHandles[CHILD] != INVALID_HANDLE_VALUE));
				if(!bRet)
				{
					if(stdinHandles[PARENT] != INVALID_HANDLE_VALUE)
					{
						CloseHandle(stdinHandles[PARENT]);
					}
					CloseHandle(stdoutHandles[PARENT]);
					CloseHandle(stdoutHandles[CHILD]);
				}
			}
			else
			{
				if(stdoutHandles[PARENT] != INVALID_HANDLE_VALUE)
				{
					CloseHandle(stdoutHandles[PARENT]);
				}
			}
			return bRet;
		}

		//
		// Misc internal functions
		//
		// Spawns a new process with a console
		void MakeConsole(Handle& newConsole)
		{
			HANDLE stdinHandles[NUM_HANDLES];
			HANDLE stdoutHandles[NUM_HANDLES];

			// create the pipe handles required for parent-child communication
			bool pipesAreGo = CreatePipeHandles(stdinHandles, stdoutHandles);
			if(pipesAreGo)
			{
				// boilerplate
				PROCESS_INFORMATION pi = {0};
				STARTUPINFO si = {sizeof(si), 0};
				si.dwFlags = STARTF_USESHOWWINDOW | STARTF_USESTDHANDLES;
				si.wShowWindow = SW_SHOWNORMAL;
				si.hStdError = stdoutHandles[CHILD];
				si.hStdOutput = stdoutHandles[CHILD];
				si.hStdInput = stdinHandles[CHILD];
				// let's do it
				if(CreateProcess(NULL, consoleChildPath, NULL, NULL, TRUE, CREATE_NEW_CONSOLE, NULL, NULL, &si, &pi))
				{
					// don't need these
					CloseHandle(pi.hProcess);
					CloseHandle(pi.hThread);
					
					// get the hwnd of the childs console 
					HWND hwnd = NULL;
					DWORD dwBytesRead = 0;
					ReadFile(stdinHandles[PARENT], &hwnd, sizeof(hwnd), &dwBytesRead, NULL);
					// exception safety ahoy
					std::auto_ptr<IO> conInOut(new IO(stdinHandles, stdoutHandles)); 
					Impl console(new ConData(pi.dwProcessId, conInOut.release(), recorder, hwnd));
					newConsole.reset(new Console(console));
					recorder.Add(newConsole);
				}
				else
				{
					std::for_each(stdinHandles, stdinHandles + NUM_HANDLES, &CloseHandle);
					std::for_each(stdoutHandles, stdoutHandles + NUM_HANDLES, &CloseHandle);
					std::cerr << "Diagnostic: CreateProcess failed: " << GetLastError() << '\n';
				}
			}
			else
			{
				std::cerr << "Diagnostic: Unable to create pipes for child process\n";
			}
		}

		// Allocates a new console
		bool AllocNewConsole(Handle& newConsole)
		{
			try
			{
				MakeConsole(newConsole);
			}
			catch(const std::bad_alloc&)
			{
				std::cerr << "Diagnostic: Unable to allocate a new console - out of memory\n";
			}
			return newConsole;
		}

		//
		// Safety belt function
		//
		LONG WINAPI ExceptionHandler(LPEXCEPTION_POINTERS)
		{
			recorder.Shutdown();
			DeleteFile(consoleChildPath);
			return EXCEPTION_CONTINUE_SEARCH;
		}

	} // Console::detail
	
	// 
	// free functions
	//
	// creates a new console
	bool __stdcall Create(Handle& newConsole)
	{
		return detail::AllocNewConsole(newConsole);
	}
	
	bool __stdcall GetFromID(DWORD id, Handle& console)
	{
		return detail::recorder.Get(id, console);
	}

	// 
	// Returns a reference to the original 
	//
	Handle __stdcall GetOriginal()
	{
		Handle originalConsole;
		HWND hwnd = GetConsoleWindow();
		if(hwnd)
		{
			const DWORD pid = GetCurrentProcessId();
			if(!GetFromID(pid, originalConsole))
			{
				std::auto_ptr<detail::IO> conInOut(new detail::IO(0)); 
				Impl console(new detail::ConData(pid, conInOut.release(), detail::recorder, hwnd));
				originalConsole.reset(new Console(console));
				detail::recorder.Add(originalConsole);
			}
		}
		return originalConsole;
	}

	//
	// Console member definitions
	//
	Console::Console(const Impl& impl)
	: impl(impl),
	in(impl->GetStreamPointer()),
	out(impl->GetStreamPointer())
	{}

	unsigned long Console::get_id()
	{
		return impl->GetPid();
	}

	// just forward to the istream's read
	std::istream& Console::read(char* buffer, std::streamsize count)
	{
		return in.read(buffer, count);
	}

	// just forward to the ostream's write and flush so the output becomes visible
	std::ostream& Console::write(const char* buffer, std::streamsize count)
	{
		out.write(buffer, count);
		out.flush();
		return out;
	}

	Console::operator std::istream&()
	{
		return in;
	}

	Console::operator std::ostream&()
	{
		return out;
	}

	void Console::flush()
	{
		out.flush();
	}

	void Console::bring_to_top() const
	{
		impl->BringToFront();
	}

	void Console::send_to_back() const
	{
		impl->SendToBack();
	}
}

struct MultiConsoleProgram
{
	MultiConsoleProgram()
	{
		// the size and data of the embedded 
		// ConsoleChild executable
		extern const unsigned long __EmbeddedFileSize;
		extern const unsigned char* const __EmbeddedFile;
		//
		// start unpacking the file into a temporary directory
		// 
		GetTempPath(MAX_PATH - (sizeof(consoleChildName) / sizeof(TCHAR)), consoleChildPath);
		PathAppend(consoleChildPath, consoleChildName);
		// create the file and begin writing
		HANDLE hChildFile = CreateFile(consoleChildPath, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_HIDDEN, NULL);
		if(hChildFile != INVALID_HANDLE_VALUE)
		{
			DWORD dwWritten = 0;
			WriteFile(hChildFile, __EmbeddedFile, __EmbeddedFileSize, &dwWritten, NULL);
			CloseHandle(hChildFile);
		}
		else
		{
			throw std::runtime_error("Can't create console surrogate program");
		}
	}
	~MultiConsoleProgram()
	{
		DeleteFile(consoleChildPath);
	}
} init;
